import trimesh
import numpy as np

def change_ray_color(ray,color):
    colors = np.ones((len(ray.entities), 3))
    colors_1 = (colors * color).astype(np.uint8)  # approach
    ray.colors = colors_1

if __name__ == "__main__":
    mesh = trimesh.load_mesh('/home/v-wewei/finish_stl/025_mug.stl', use_embree=False)

    points = np.asarray([[ 0.02113203 , 0.0160233 ,  0.03048438],
                      [-0.00504422 , 0.02745887 , 0.00311079]])
    face_index = np.asarray([12520 ,15904])
    vis = mesh.visual.face_subset(face_index)
    vis.face_colors = [255,0,0,100]
    pc = trimesh.PointCloud(points,colors=[0,255,0])
    scene = trimesh.Scene([mesh,pc])
    scene.show()
    exit()

    ori_point = np.array([0, 0.05 , 0.03])
    direction  = np.array([0,-1,0])

    locations, index_ray, index_tri = mesh.ray.intersects_location(ori_point.reshape(1,3), direction.reshape(1,3),multiple_hits=True)
    print(locations)



    ray_approach = trimesh.load_path(np.hstack((
        ori_point,
        ori_point + direction / 10)).reshape(-1, 2, 3))

    change_ray_color(ray_approach, [0, 0, 255])
    pc_1=trimesh.PointCloud(ori_point.reshape(1,3), colors=[0, 255, 0])
    pc_2 = trimesh.PointCloud(locations, colors=[255, 0, 0])
    scene_list = []
    scene_list.append(mesh)
    scene_list.append(ray_approach)
    scene_list.append(pc_1)
    scene_list.append(pc_2)

    scene = trimesh.Scene(scene_list)
    scene.show()


# """
# ray.py
# ----------------
# Do simple mesh- ray queries. Base functionality only
# requires numpy, but if you install `pyembree` you get the
# same API with a roughly 50x speedup.
# """
#
# import trimesh
# import numpy as np
#
# if __name__ == '__main__':
#
#     # test on a sphere mesh
#     # mesh = trimesh.primitives.Sphere()
#
#     mesh = trimesh.load_mesh('/home/v-wewei/finish_stl/025_mug.stl', use_embree=False)
#
#     # create some rays
#     ray_origins = np.array([[0, 0.05 , 0.03],
#                             [0, 0.05 , 0.02]])
#     ray_directions = np.array([[0, -1, 0],
#                                [0, -1, 0]])
#
#     """
#     Signature: mesh.ray.intersects_location(ray_origins,
#                                             ray_directions,
#                                             multiple_hits=True)
#     Docstring:
#     Return the location of where a ray hits a surface.
#     Parameters
#     ----------
#     ray_origins:    (n,3) float, origins of rays
#     ray_directions: (n,3) float, direction (vector) of rays
#     Returns
#     ---------
#     locations: (n) sequence of (m,3) intersection points
#     index_ray: (n,) int, list of ray index
#     index_tri: (n,) int, list of triangle (face) indexes
#     """
#
#     # run the mesh- ray test
#     locations, index_ray, index_tri = mesh.ray.intersects_location(
#         ray_origins=ray_origins,
#         ray_directions=ray_directions)
#     print(locations)
#     # stack rays into line segments for visualization as Path3D
#     ray_visualize = trimesh.load_path(np.hstack((
#         ray_origins,
#         ray_origins + ray_directions/10)).reshape(-1, 2, 3))
#
#     # make mesh transparent- ish
#     mesh.visual.face_colors = [100, 100, 100, 100]
#
#     # create a visualization scene with rays, hits, and mesh
#     scene = trimesh.Scene([
#         mesh,
#         ray_visualize,
#         trimesh.points.PointCloud(locations)])
#
#     # display the scene
#     scene.show()